import os
import base64
import re
from PIL import Image, ImageDraw
import dashscope

# -------------------------- 1. 配置API密钥 --------------------------
DASHSCOPE_API_KEY = "sk-be24ebb43e41447ea337f5dbf17617c4"
dashscope.api_key = DASHSCOPE_API_KEY


# -------------------------- 2. 增强型推理函数（强制坐标输出） --------------------------
def inference_local_image(img_local_path, prompt):
    # 1. 基础校验
    if not os.path.exists(img_local_path):
        raise Exception(f"图片不存在：{img_local_path}\n请检查路径中的'x'是否为实际用户名，且output.png存在")
    valid_formats = (".png", ".jpg", ".jpeg")
    if not img_local_path.lower().endswith(valid_formats):
        raise Exception(f"仅支持{valid_formats}格式，当前格式：{img_local_path.split('.')[-1]}")

    # 2. 图片转base64（含MIME类型）
    try:
        img_format = img_local_path.split(".")[-1].lower()
        mime_type = f"image/{img_format}" if img_format != "jpg" else "image/jpeg"
        with open(img_local_path, "rb") as f:
            img_base64 = base64.b64encode(f.read()).decode("utf-8")
        img_url = f"data:{mime_type};base64,{img_base64}"
        if len(img_url) < 1000:
            raise Exception("图片损坏，base64编码长度异常")
    except Exception as e:
        raise Exception(f"图片转base64失败：{str(e)}")

    # 3. 构造API请求（关键：添加坐标格式示例，降低模型理解成本）
    messages = [
        {
            "role": "user",
            "content": [
                {"image": img_url},
                # 增强prompt：明确坐标格式+示例，强制模型输出
                prompt + "\n坐标格式要求：每个题目、选项、表格都必须标注(x1,y1,x2,y2)，示例：\n### 一、选择题 (10,20,800,100)\n1. 题目内容 (20,30,780,60)\nA. 选项 (30,70,180,90)  B. 选项 (200,70,350,90)..."
            ]
        }
    ]

    # 4. 调用API（用qwen-vl-max，坐标识别优先级更高）
    try:
        print("调用Qwen3-VL API中...（约10-25秒，勿中断）")
        response = dashscope.MultiModalConversation.call(
            model="qwen-vl-max",  # 必换max模型，plus对纯文本坐标支持弱
            messages=messages,
            stream=False,
            timeout=60  # 延长超时，确保复杂标注完成
        )
    except Exception as e:
        raise Exception(f"API调用异常：{str(e)}")

    # 5. 解析结果
    if response.status_code != 200:
        raise Exception(f"API错误：{response.message}（状态码：{response.status_code}）")
    content = response.output["choices"][0]["message"]["content"]
    return content[0] if isinstance(content, list) else content


# -------------------------- 3. 坐标标注函数（不变） --------------------------
def draw_bbox_on_local_img(img_local_path, markdown_text, save_path="marked_img.png"):
    try:
        img = Image.open(img_local_path)
        draw = ImageDraw.Draw(img)
    except Exception as e:
        raise Exception(f"打开标注图片失败：{str(e)}")

    bbox_pattern = r"\((\d+),\s*(\d+),\s*(\d+),\s*(\d+)\)"
    bboxes = re.findall(bbox_pattern, str(markdown_text))

    if not bboxes:
        print("⚠️  未检测到坐标，保存原图")
        img.save(save_path)
        return

    border_color = "#FF0000"
    border_width = 3
    for bbox in bboxes:
        x1, y1, x2, y2 = map(int, bbox)
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(img.width, x2), min(img.height, y2)
        draw.rectangle([x1, y1, x2, y2], outline=border_color, width=border_width)

    img.save(save_path)
    print(f"✅ 标注图保存至：{os.path.abspath(save_path)}")


# -------------------------- 4. 主逻辑 --------------------------
if __name__ == "__main__":
    LOCAL_IMG_PATH = r"C:\Users\x\Desktop\QWen3识别文档\output.png"  # 替换x为实际用户名
    # 精简且强制的prompt
    PROMPT = "解析图片为QwenVL格式Markdown，表格用LaTeX，每个题目、选项、表格必须标注精确(x1,y1,x2,y2)坐标，不遗漏任何元素"

    try:
        print("=" * 60)
        print("任务：本地图片解析 -> 含坐标Markdown -> 可视化标注")
        print(f"处理图片：{LOCAL_IMG_PATH}")
        print("=" * 60)

        markdown_result = inference_local_image(LOCAL_IMG_PATH, PROMPT)
        
        print("\n" + "=" * 80)
        print("生成的QwenVL Markdown（含坐标）：")
        print("-" * 80)
        print(markdown_result)
        print("=" * 80)

        draw_bbox_on_local_img(LOCAL_IMG_PATH, markdown_result)
        print("\n🎉 所有任务完成！")

    except Exception as e:
        print(f"\n❌ 执行失败：{str(e)}")
        print("\n排查重点：1. 更换为带图表的图片；2. 确认API密钥有max模型调用额度；3. 关闭代理确保网络通畅。")